from django.test import TestCase
from django.urls import reverse
from rest_framework.test import APITestCase, APIClient
from rest_framework import status
from .models import Industry, KeywordTrend, PatentDistribution, ProvincePatent
from users.models import User
import os
from django.conf import settings
import random

class TrendAnalysisTests(APITestCase):
    def setUp(self):
        # 创建测试用户
        self.user = User.objects.create_user(
            username='testuser',
            password='testpass123'
        )
        self.client = APIClient()
        self.client.force_authenticate(user=self.user)
        
        # 创建测试数据
        self.industry = Industry.objects.create(
            id=1,
            name='测试产业',
            description='测试描述'
        )
        
        # 创建关键词趋势数据
        self.keyword = '测试关键词'
        for year in range(2019, 2024):
            KeywordTrend.objects.create(
                industry=self.industry,
                keyword=self.keyword,
                year=year,
                count=100 + year - 2019
            )
            
        # 创建专利分布数据
        self.institution = '测试机构'
        for year in range(2019, 2024):
            PatentDistribution.objects.create(
                industry=self.industry,
                institution=self.institution,
                year=year,
                patent_count=150
            )

    def test_get_industry_list(self):
        """测试获取产业列表"""
        url = reverse('industry-list')
        response = self.client.get(url)
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(len(response.data), 1)
        self.assertEqual(response.data[0]['name'], '测试产业')

    def test_get_industry_trend(self):
        """测试获取产业趋势分析"""
        url = reverse('industry-trend', args=[self.industry.id])
        response = self.client.get(url)
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        
        data = response.data['data']
        self.assertEqual(data['industryName'], '测试产业')
        self.assertEqual(len(data['keywordTrends']['years']), 5)
        self.assertEqual(len(data['patentDistribution']), 1) 

    def test_get_province_patents(self):
        """测试获取省份专利分布"""
        # 创建测试数据
        provinces = ['北京', '上海', '广东']
        years = range(2019, 2024)
        
        for province in provinces:
            for year in years:
                ProvincePatent.objects.create(
                    industry=self.industry,
                    province=province,
                    year=year,
                    patent_count=100
                )
        
        # 测试获取所有年份数据
        url = reverse('industry-province-patents', args=[self.industry.id])
        response = self.client.get(url)
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        
        data = response.data['data']
        self.assertEqual(len(data['years']), 5)
        self.assertEqual(len(data['provinces']), 3)
        self.assertEqual(len(data['data']), 3)
        
        # 测试获取特定年份数据
        url = f"{url}?year=2020"
        response = self.client.get(url)
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        
        data = response.data['data']
        self.assertEqual(len(data['years']), 1)
        self.assertEqual(data['years'][0], 2020) 

    def test_generate_patent_heatmap(self):
        """测试生成专利热力图"""
        # 创建测试数据
        provinces = [
            '北京', '上海', '广东', '江苏', '浙江',
            '四川', '湖北', '陕西', '山东', '河南'
        ]
        year = 2023
        
        for province in provinces:
            ProvincePatent.objects.create(
                industry=self.industry,
                province=province,
                year=year,
                patent_count=random.randint(100, 1000)
            )
        
        # 测试生成热力图
        url = reverse('industry-patent-heatmap', args=[self.industry.id])
        response = self.client.get(f'{url}?year={year}')
        
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(response.data['data']['year'], str(year))
        
        # 验证图片是否生成
        image_url = response.data['data']['imageUrl']
        image_path = os.path.join(settings.MEDIA_ROOT, image_url.lstrip('/'))
        self.assertTrue(os.path.exists(image_path)) 